package com.wang.sort3;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.PriorityQueue;
import java.util.Set;
import java.util.Stack;
import com.wang.sort3.Demo08_GraphGenerator.*;

/**
 * @Author: along
 * @Create: 2021/4/24
 * @Description:最小生成树算法Kruskal
 */
public class Demo12_Kruskal {
    //undirected graph only
        public static class MySets{
            public HashMap<Node, List<Node>>  setMap;
            public MySets(List<Node> nodes) {
                for(Node cur : nodes) {
                    List<Node> set = new ArrayList<Node>();
                    set.add(cur);
                    setMap.put(cur, set);
                }
            }

            public boolean isSameSet(Node from, Node to) {
                List<Node> fromSet  = setMap.get(from);
                List<Node> toSet = setMap.get(to);
                return fromSet == toSet;
            }

            public void union(Node from, Node to) {
                List<Node> fromSet  = setMap.get(from);
                List<Node> toSet = setMap.get(to);
                for(Node toNode : toSet) {
                    fromSet.add(toNode);
                    setMap.put(toNode, fromSet);
                }
            }
        }

        // Union-Find Set
        public static class UnionFind {
            // key 某一个节点， value key节点往上的节点
            private HashMap<Node, Node> fatherMap;
            // key 某一个集合的代表节点, value key所在集合的节点个数
            private HashMap<Node, Integer> sizeMap;

            public UnionFind() {
                fatherMap = new HashMap<Node, Node>();
                sizeMap = new HashMap<Node, Integer>();
            }

            public void makeSets(Collection<Node> nodes) {
                fatherMap.clear();
                sizeMap.clear();
                for (Node node : nodes) {
                    fatherMap.put(node, node);
                    sizeMap.put(node, 1);
                }
            }

            private Node findFather(Node n) {
                Stack<Node> path = new Stack<>();
                while(n != fatherMap.get(n)) {
                    path.add(n);
                    n = fatherMap.get(n);
                }
                while(!path.isEmpty()) {
                    fatherMap.put(path.pop(), n);
                }
                return n;
            }

            public boolean isSameSet(Node a, Node b) {
                return findFather(a) == findFather(b);
            }

            public void union(Node a, Node b) {
                if (a == null || b == null) {
                    return;
                }
                Node aDai = findFather(a);
                Node bDai = findFather(b);
                if (aDai != bDai) {
                    int aSetSize = sizeMap.get(aDai);
                    int bSetSize = sizeMap.get(bDai);
                    if (aSetSize <= bSetSize) {
                        fatherMap.put(aDai, bDai);
                        sizeMap.put(bDai, aSetSize + bSetSize);
                        sizeMap.remove(aDai);
                    } else {
                        fatherMap.put(bDai, aDai);
                        sizeMap.put(aDai, aSetSize + bSetSize);
                        sizeMap.remove(bDai);
                    }
                }
            }
        }


        public static class EdgeComparator implements Comparator<Edge> {

            @Override
            public int compare(Edge o1, Edge o2) {
                return o1.weight - o2.weight;
            }

        }

        public static Set<Edge> kruskalMST(Graph graph) {
            UnionFind unionFind = new UnionFind();
            unionFind.makeSets(graph.nodes.values());
            PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new EdgeComparator());
            //把所有的边加入堆中，按照自己制定的策略(权重小的在堆顶)
            for (Edge edge : graph.edges) { // M 条边
                priorityQueue.add(edge);  // O(logM)
            }
            Set<Edge> result = new HashSet<>();//结果集
            while (!priorityQueue.isEmpty()) { // M 条边
                Edge edge = priorityQueue.poll(); // O(logM)
                //如果边的from和to不属于同一个集合，那么就需要这个边，并且联合from和to
                if (!unionFind.isSameSet(edge.from, edge.to)) { // O(1)
                    result.add(edge);
                    unionFind.union(edge.from, edge.to);
                }
            }
            return result;
        }

}
